{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2a5b47fb-b49f-42b1-8c12-09e8c963c921",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "数据集大小: 15011\n",
      "\n",
      "随机数据样例:\n",
      "{'instruction': 'Which of the following can fly: bird, cow, cat, dog, car, helicopter, drone, bike, pot, butterfly, feather?', 'context': '', 'response': \"The objects that can fly are: bird, helicopter, drone, butterfly.\\nThe objects that cannot fly are: cow, cat, dog, car, bike, pot, feather.\\n\\nHowever, someone may claim that a feather can fly as well because it's light enough and a mild wind blow could make it fly.\", 'category': 'classification'}\n",
      "\n",
      "格式化后的数据样例:\n",
      "### Instruction:\n",
      "Use the Input below to create an instruction, which could have been used to generate the input using an LLM. \n",
      "\n",
      "### Input:\n",
      "The Tampa Bay Rowdies wear green and yellow and play in Al Lang Stadium in downtown St. Petersburg.\n",
      "\n",
      "### Response:\n",
      "Based on the reference text, what colors do the Tampa Bay Rowdies wear and what stadium do they play in?\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/hlf_old_env/lib/python3.10/site-packages/huggingface_hub/file_download.py:945: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "efe83081417a45f8abeaaaa4e198664e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 8,388,608 || all params: 6,746,804,224 || trainable%: 0.12433454005023165\n",
      "\n",
      "开始训练...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/hlf_old_env/lib/python3.10/site-packages/trl/trainer/sft_trainer.py:341: UserWarning: You passed `packing=True` to the SFTTrainer, and you are training your model with `max_steps` strategy. The dataset will be iterated until the `max_steps` are reached.\n",
      "  warnings.warn(\n",
      "/root/miniconda3/envs/hlf_old_env/lib/python3.10/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='100' max='100' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [100/100 18:10, Epoch 0/1]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>20</td>\n",
       "      <td>2.148100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>40</td>\n",
       "      <td>1.625300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>60</td>\n",
       "      <td>1.519600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>80</td>\n",
       "      <td>1.467900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>100</td>\n",
       "      <td>1.476400</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/hlf_old_env/lib/python3.10/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/root/miniconda3/envs/hlf_old_env/lib/python3.10/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/root/miniconda3/envs/hlf_old_env/lib/python3.10/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/root/miniconda3/envs/hlf_old_env/lib/python3.10/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "保存模型到: models/llama-7-int4-dolly-20250816_113249\n",
      "训练完成!\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments\n",
    "from peft import LoraConfig, prepare_model_for_kbit_training, get_peft_model\n",
    "from trl import SFTTrainer\n",
    "import datetime\n",
    "from datasets import load_dataset\n",
    "from random import randrange\n",
    "\n",
    "# 从Hugging Face Hub加载databricks-dolly-15k数据集\n",
    "# 该数据集包含15,000条人工生成的指令-响应对，用于训练对话AI\n",
    "dataset = load_dataset(\"databricks/databricks-dolly-15k\", split=\"train\")\n",
    "print(\"数据集大小:\", len(dataset))\n",
    "# 随机选择并打印一条数据样例，了解数据结构\n",
    "sample = dataset[randrange(len(dataset))]\n",
    "print(\"\\n随机数据样例:\")\n",
    "print(sample)\n",
    "\n",
    "def format_instruction(sample_data):\n",
    "    \"\"\"\n",
    "    将数据格式化为指令微调的标准格式（Alpaca格式）\n",
    "    \n",
    "    参数:\n",
    "        sample_data (dict): 包含'response'和'instruction'键的字典\n",
    "        \n",
    "    返回:\n",
    "        str: 格式化后的字符串，包含指令、输入和响应三部分\n",
    "        \n",
    "    格式示例:\n",
    "        ### Instruction:\n",
    "        Use the Input below to create an instruction...\n",
    "        \n",
    "        ### Input:\n",
    "        {response}\n",
    "        \n",
    "        ### Response:\n",
    "        {instruction}\n",
    "    \"\"\"\n",
    "    # 检查数据是否包含所需字段\n",
    "    if 'response' not in sample_data or 'instruction' not in sample_data:\n",
    "        return \"错误: 数据中缺少'response'或'instruction'字段\"\n",
    "\n",
    "    return f\"\"\"### Instruction:\n",
    "Use the Input below to create an instruction, which could have been used to generate the input using an LLM. \n",
    "\n",
    "### Input:\n",
    "{sample_data['response']}\n",
    "\n",
    "### Response:\n",
    "{sample_data['instruction']}\n",
    "\"\"\"\n",
    "\n",
    "# 打印格式化后的数据样例\n",
    "print(\"\\n格式化后的数据样例:\")\n",
    "print(format_instruction(dataset[randrange(len(dataset))]))\n",
    "\n",
    "# 创建时间戳用于保存模型\n",
    "timestamp = datetime.datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
    "\n",
    "# 训练配置参数\n",
    "output_dir = f\"models/llama-7-int4-dolly-{timestamp}\"  # 模型输出目录\n",
    "\n",
    "# Flash Attention设置（需要单独安装且硬件支持，已测试过，当前硬件不支持）\n",
    "use_flash_attention = False  # 默认关闭，安装flash-attn后可设为True\n",
    "\n",
    "# 模型选择（使用不需要Meta审核的版本）\n",
    "model_id = \"NousResearch/Llama-2-7b-hf\" \n",
    "\n",
    "# 4-bit量化配置（使用BitsAndBytes进行高效加载）\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,  # 启用4-bit量化\n",
    "    bnb_4bit_use_double_quant=True,  # 嵌套量化，进一步减少内存占用\n",
    "    bnb_4bit_quant_type=\"nf4\",  # 4-bit量化类型（NF4最优）\n",
    "    bnb_4bit_compute_dtype=torch.float16 # 计算时使用float16\n",
    ")\n",
    "\n",
    "# 加载量化后的模型\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_id,\n",
    "    quantization_config=bnb_config,  # 应用量化配置\n",
    "    use_cache=False,  # 禁用缓存（与梯度检查点冲突）\n",
    "    device_map=\"auto\"  # 自动分配设备（GPU/CPU）\n",
    ")\n",
    "model.config.pretraining_tp = 1  # 设置张量并行参数\n",
    "\n",
    "# 验证是否成功启用Flash Attention\n",
    "if use_flash_attention:\n",
    "    from utils.llama_patch import forward    \n",
    "    assert model.model.layers[0].self_attn.forward.__doc__ == forward.__doc__, \"Flash Attention未启用\"\n",
    "\n",
    "# 加载分词器\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "tokenizer.pad_token = tokenizer.eos_token  # 使用EOS token作为填充token\n",
    "tokenizer.padding_side = \"right\"  # 填充方向（右侧）\n",
    "\n",
    "# LoRA配置（参数高效微调）\n",
    "peft_config = LoraConfig(\n",
    "    lora_alpha=16,  # LoRA缩放因子\n",
    "    lora_dropout=0.1,  # Dropout率\n",
    "    r=16,  # LoRA秩（矩阵分解维度）\n",
    "    bias=\"none\",  # 不使用偏置项\n",
    "    task_type=\"CAUSAL_LM\",  # 任务类型（因果语言模型）\n",
    ")\n",
    "\n",
    "# 准备模型进行k-bit训练\n",
    "model = prepare_model_for_kbit_training(model)\n",
    "# 应用QLoRA配置\n",
    "qlora_model = get_peft_model(model, peft_config)\n",
    "# 打印可训练参数数量（仅LoRA参数会被训练）\n",
    "qlora_model.print_trainable_parameters()\n",
    "\n",
    "# 训练参数配置\n",
    "args = TrainingArguments(\n",
    "    output_dir=output_dir,  # 输出目录\n",
    "    num_train_epochs=3,  # 训练轮次\n",
    "    max_steps=100,  # 最大训练步数\n",
    "    per_device_train_batch_size=1,  # 每个设备的batch大小（根据显存调整）\n",
    "    gradient_accumulation_steps=8,  # 梯度累积步数，通过梯度累积模拟更大batch\n",
    "    gradient_checkpointing=True,  # 梯度检查点（节省显存）\n",
    "    optim=\"paged_adamw_8bit\",  # 分页优化器（防止OOM，使用8bit优化器（比32bit省显存））\n",
    "    logging_steps=20,  # 日志记录间隔\n",
    "    save_strategy=\"steps\",  # 保存策略（按步数保存）\n",
    "    save_steps=20,  # 保存间隔（每20步保存一次）\n",
    "    learning_rate=2e-4,  # 学习率\n",
    "    fp16=True,  # 用fp16替代bf16（T4不支持bf16加速）\n",
    "    max_grad_norm=0.3,  # 梯度裁剪阈值\n",
    "    warmup_ratio=0.03,  # 预热比例\n",
    "    lr_scheduler_type=\"constant\"  # 学习率调度器类型\n",
    ")\n",
    "\n",
    "# 创建监督式微调训练器\n",
    "trainer = SFTTrainer(\n",
    "    model=qlora_model,  # QLoRA模型\n",
    "    train_dataset=dataset,  # 训练数据集\n",
    "    peft_config=peft_config,  # PEFT配置\n",
    "    max_seq_length=128,  # 最大序列长度\n",
    "    tokenizer=tokenizer,  # 分词器\n",
    "    packing=True,  # 序列打包（提高效率）\n",
    "    formatting_func=format_instruction,  # 数据格式化函数\n",
    "    args=args,  # 训练参数\n",
    ")\n",
    "\n",
    "# 开始训练\n",
    "print(\"\\n开始训练...\")\n",
    "trainer.train()\n",
    "\n",
    "# 保存训练后的模型\n",
    "print(\"\\n保存模型到:\", output_dir)\n",
    "trainer.save_model()\n",
    "print(\"训练完成!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "251a542f-8535-408b-935b-11eedfdb3550",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/hlf_old_env/lib/python3.10/site-packages/huggingface_hub/file_download.py:945: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "26b8869ea6b44fbf9620957605dd262a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.\n",
      "/root/miniconda3/envs/hlf_old_env/lib/python3.10/site-packages/bitsandbytes/nn/modules.py:226: UserWarning: Input type into Linear4bit is torch.float16, but bnb_4bit_compute_type=torch.float32 (default). This will lead to slow inference or training speed.\n",
      "  warnings.warn(f'Input type into Linear4bit is torch.float16, but bnb_4bit_compute_type=torch.float32 (default). This will lead to slow inference or training speed.')\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "===== 输入与生成结果对比 =====\n",
      "\n",
      "[原始输入]:\n",
      "Marshan Lynch\n",
      "\n",
      "[模型生成的指令]:\n",
      "What is Marshan Lynch's last name?\n",
      "\n",
      "[真实指令]:\n",
      "What athlete created the 'beast quake' for the Seattle Seahawks?\n",
      "\n",
      "===== 执行说明 =====\n",
      "1. 模型从以下目录加载: models/llama-7-int4-dolly-20250816_113249\n",
      "2. 当前使用的设备: GPU\n",
      "3. 输入长度: 51 tokens\n",
      "4. 输出长度: 64 tokens\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from peft import AutoPeftModelForCausalLM\n",
    "from transformers import AutoTokenizer\n",
    "from datasets import load_dataset \n",
    "from random import randrange\n",
    "\n",
    "# 模型目录路径（替换为您实际训练保存的模型路径）\n",
    "model_dir = \"models/llama-7-int4-dolly-20250816_113249\"\n",
    " \n",
    "# 加载经过QLoRA微调的模型\n",
    "# 参数说明：\n",
    "# - low_cpu_mem_usage: 减少CPU内存占用\n",
    "# - torch_dtype: 使用float16精度加载\n",
    "# - load_in_4bit: 启用4-bit量化以节省显存\n",
    "model = AutoPeftModelForCausalLM.from_pretrained(\n",
    "    model_dir,\n",
    "    low_cpu_mem_usage=True,\n",
    "    torch_dtype=torch.float16,\n",
    "    load_in_4bit=True,\n",
    ") \n",
    "\n",
    "# 加载与模型匹配的分词器\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_dir)\n",
    "\n",
    "# 从Hugging Face Hub加载databricks-dolly-15k数据集\n",
    "# 这是一个包含15,000条人工生成的指令-响应对的数据集\n",
    "dataset = load_dataset(\"databricks/databricks-dolly-15k\", split=\"train\")\n",
    "\n",
    "# 随机选择一条数据样例\n",
    "sample = dataset[randrange(len(dataset))]\n",
    " \n",
    "# 构建推理提示模板\n",
    "# 格式说明：\n",
    "# - Instruction: 任务描述\n",
    "# - Input: 模型需要处理的输入内容\n",
    "# - Response: 留空等待模型生成\n",
    "prompt = f\"\"\"### Instruction:\n",
    "Use the Input below to create an instruction, which could have been used to generate the input using an LLM. \n",
    " \n",
    "### Input:\n",
    "{sample['response']}\n",
    " \n",
    "### Response:\n",
    "\"\"\"\n",
    " \n",
    "# 使用分词器处理提示文本：\n",
    "# - return_tensors=\"pt\": 返回PyTorch张量\n",
    "# - truncation=True: 超过最大长度时自动截断\n",
    "# - .cuda(): 将输入转移到GPU（如果可用）\n",
    "input_ids = tokenizer(prompt, return_tensors=\"pt\", truncation=True).input_ids.cuda()\n",
    "\n",
    "# 模型推理生成\n",
    "# 参数说明：\n",
    "# - max_new_tokens: 最大生成token数\n",
    "# - do_sample: 启用随机采样（非贪婪解码）\n",
    "# - top_p=0.9: 核采样概率阈值\n",
    "# - temperature=0.9: 控制生成随机性的温度参数\n",
    "outputs = model.generate(\n",
    "    input_ids=input_ids, \n",
    "    max_new_tokens=100, \n",
    "    do_sample=True, \n",
    "    top_p=0.9,\n",
    "    temperature=0.9\n",
    ")\n",
    "\n",
    "# 打印结果对比\n",
    "print(\"===== 输入与生成结果对比 =====\")\n",
    "print(f\"\\n[原始输入]:\\n{sample['response']}\\n\")\n",
    "print(f\"[模型生成的指令]:\\n{tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0][len(prompt):]}\")\n",
    "print(f\"[真实指令]:\\n{sample['instruction']}\")\n",
    "\n",
    "# 输出信息说明\n",
    "print(\"\\n===== 执行说明 =====\")\n",
    "print(\"1. 模型从以下目录加载:\", model_dir)\n",
    "print(\"2. 当前使用的设备:\", \"GPU\" if torch.cuda.is_available() else \"CPU\")\n",
    "print(\"3. 输入长度:\", len(input_ids[0]), \"tokens\")\n",
    "print(\"4. 输出长度:\", len(outputs[0]), \"tokens\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6882c96-9cee-4960-a7cd-6aa5ccf52183",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "hlf_old_env",
   "language": "python",
   "name": "hlf_old_env"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
